# File: src/prepare_data.py

import os
import json
import pandas as pd

def load_json_with_fallback(path):
    """
    尝试使用utf-8读取JSON，若失败则使用gbk。
    """
    try:
        with open(path, 'r', encoding='utf-8') as f:
            return json.load(f)
    except UnicodeDecodeError:
        with open(path, 'r', encoding='gbk', errors='ignore') as f:
            return json.load(f)

def prepare_hrv_dataset(data_folder, output_csv, gender_limit='无'):
    """
    读取 data_folder 下的所有 .mat 对应的 .json + _hrv.json 文件，
    按照筛选条件合并到一个 DataFrame，再输出到 CSV。

    筛选条件：
      1. AnalysisText 非空
      2. nbeats > 40000
      3. 0 < age < 100
      4. AnalysisText 中不含 ['颤','停','就医','频']
      5. 如果 gender_limit 为 '男' 或 '女'，则只保留该性别；'无' 则保留所有
    """
    records = []
    for fname in os.listdir(data_folder):
        if not fname.endswith('.mat'):
            continue
        base = fname[:-4]  # 如 "138022200000"
        path_user = os.path.join(data_folder, base + '.json')
        path_hrv  = os.path.join(data_folder, base + '_hrv.json')
        if not (os.path.exists(path_user) and os.path.exists(path_hrv)):
            continue

        # 读两份 JSON（带编码回退）
        user = load_json_with_fallback(path_user)
        hrv  = load_json_with_fallback(path_hrv)

        age      = user.get('age')
        gender   = user.get('Gender')
        analysis = user.get('AnalysisText', '')

        # 性别筛选
        if gender_limit in ('男','女') and gender != gender_limit:
            continue

        # 基本筛选
        if (analysis
            and hrv.get('nbeats', 0) > 40000
            and isinstance(age, (int, float)) and 0 < age < 100
            and all(term not in analysis for term in ['颤','停','就医','频'])):

            row = {
                'record': base,
                'age': age,
                'Gender': gender,
                'AnalysisText': analysis
            }
            row.update(hrv)
            records.append(row)

    # 合并 DataFrame 并保存
    df = pd.DataFrame(records)
    os.makedirs(os.path.dirname(output_csv), exist_ok=True)
    df.to_csv(output_csv, index=False, encoding='utf-8')
    print(f"已生成 {len(df)} 条记录，保存到 {output_csv}")

if __name__ == '__main__':
    data_folder = r"C:\py\hrv_predict\data_hrv"
    output_csv  = r"C:\py\hrv_predict\data_hrv\hrv_dataset.csv"
    prepare_hrv_dataset(data_folder, output_csv, gender_limit='无')
